package com.techzhi.tidb.pool;

import com.techzhi.tidb.config.TiDBConfigManager;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import com.zaxxer.hikari.HikariPoolMXBean;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import javax.annotation.PreDestroy;
import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * TiDB连接池管理器
 * 专门针对TiDB优化的连接池，考虑性能和安全
 *
 * @author shouzhi
 * @date 2025-09-01
 */

@Slf4j
public class TiDBConnectionPoolManager {

    private volatile HikariDataSource dataSource;
    private final Object lock = new Object();
    private final AtomicBoolean initialized = new AtomicBoolean(false);

    @Autowired
    private TiDBConfigManager configManager;

    /**
     * 获取连接池（懒加载）
     */
    public DataSource getDataSource() {
        if (!initialized.get() || dataSource == null || dataSource.isClosed()) {
            synchronized (lock) {
                if (!initialized.get() || dataSource == null || dataSource.isClosed()) {
                    initializeDataSource();
                }
            }
        }
        return dataSource;
    }

    /**
     * 获取连接
     */
    public Connection getConnection() throws SQLException {
        return getDataSource().getConnection();
    }

    /**
     * 初始化数据源
     */
    private void initializeDataSource() {
        if (!configManager.isEnabled()) {
            throw new IllegalStateException("TiDB未启用，请检查配置 tidb.enabled=true");
        }

        if (!configManager.validateConfig()) {
            throw new IllegalStateException("TiDB配置无效，请检查配置参数");
        }

        try {
            // 关闭旧的数据源
            if (dataSource != null && !dataSource.isClosed()) {
                dataSource.close();
            }

            HikariConfig hikariConfig = createHikariConfig();
            this.dataSource = new HikariDataSource(hikariConfig);

            // 预热连接池
            warmUpConnectionPool();

            initialized.set(true);
            log.info("TiDB连接池初始化成功: {}", configManager.getMaskedConnectionInfo());

        } catch (Exception e) {
            initialized.set(false);
            log.error("TiDB连接池初始化失败: {}", configManager.getMaskedConnectionInfo(), e);
            throw new RuntimeException("TiDB连接池初始化失败", e);
        }
    }

    /**
     * 创建Hikari配置（针对TiDB优化）
     */
    private HikariConfig createHikariConfig() {
        HikariConfig config = new HikariConfig();
        TiDBConfigManager.TiDBConfig tidbConfig = configManager.getConfig();
        TiDBConfigManager.PoolConfig poolConfig = configManager.getPool();
        TiDBConfigManager.SecurityConfig securityConfig = configManager.getSecurity();

        // 基础连接配置
        config.setJdbcUrl(configManager.buildJdbcUrl());
        config.setUsername(tidbConfig.getUsername());
        config.setPassword(configManager.getDecryptedPassword());
        config.setDriverClassName("com.mysql.cj.jdbc.Driver");

        // 连接池配置
        config.setMaximumPoolSize(poolConfig.getMaximumPoolSize());
        config.setMinimumIdle(poolConfig.getMinimumIdle());
        config.setConnectionTimeout(poolConfig.getConnectionTimeout());
        config.setIdleTimeout(poolConfig.getIdleTimeout());
        config.setMaxLifetime(poolConfig.getMaxLifetime());
        config.setPoolName(poolConfig.getPoolName());
        config.setAutoCommit(poolConfig.isAutoCommit());

        // 安全配置
        config.setConnectionTestQuery(securityConfig.getConnectionTestQuery());
        config.setValidationTimeout(securityConfig.getValidationTimeout());
        config.setLeakDetectionThreshold(poolConfig.getLeakDetectionThreshold());

        // TiDB性能优化配置
        optimizeForTiDB(config);

        return config;
    }

    /**
     * TiDB性能优化配置
     */
    private void optimizeForTiDB(HikariConfig config) {
        // MySQL兼容性优化
        config.addDataSourceProperty("cachePrepStmts", "true");
        config.addDataSourceProperty("prepStmtCacheSize", "250");
        config.addDataSourceProperty("prepStmtCacheSqlLimit", "2048");
        config.addDataSourceProperty("useServerPrepStmts", "true");

        // TiDB批量操作优化
        config.addDataSourceProperty("rewriteBatchedStatements", "true");
        config.addDataSourceProperty("allowMultiQueries", "true");

        // 网络优化
        config.addDataSourceProperty("useCompression", "true");
        config.addDataSourceProperty("tcpKeepAlive", "true");
        config.addDataSourceProperty("tcpNoDelay", "true");

        // 结果集优化
        config.addDataSourceProperty("cacheResultSetMetadata", "true");
        config.addDataSourceProperty("cacheServerConfiguration", "true");
        config.addDataSourceProperty("useLocalSessionState", "true");

        // 事务优化
        config.addDataSourceProperty("elideSetAutoCommits", "true");
        config.addDataSourceProperty("maintainTimeStats", "false");

        // TiDB特定优化
        config.addDataSourceProperty("useAffectedRows", "true");
        config.addDataSourceProperty("useLocalTransactionState", "true");

        log.debug("TiDB性能优化配置已应用");
    }

    /**
     * 预热连接池
     * <p>
     * 预热连接池的目的是在应用启动时提前创建并验证一定数量的数据库连接，
     * 避免在首次请求时因建立连接而导致的延迟，提高系统响应速度和用户体验。
     * 同时也可以提前发现连接配置是否正确，避免运行时才发现连接问题。
     */
    private void warmUpConnectionPool() {
        try (Connection conn = dataSource.getConnection()) {
            // 预热操作 SELECT 1
            conn.prepareStatement(configManager.getSecurity().getConnectionTestQuery()).execute();
            log.info("TiDB连接池预热完成");
        } catch (Exception e) {
            log.warn("TiDB连接池预热失败，但不影响正常使用", e);
        }
    }

    /**
     * 获取连接池统计信息
     */
    public Map<String, Object> getPoolStats() {
        if (!initialized.get() || dataSource == null) {
            return Collections.emptyMap();
        }

        try {
            HikariPoolMXBean poolBean = dataSource.getHikariPoolMXBean();
            Map<String, Object> stats = new HashMap<>();
            stats.put("poolName", configManager.getPool().getPoolName());
            stats.put("activeConnections", poolBean.getActiveConnections());
            stats.put("idleConnections", poolBean.getIdleConnections());
            stats.put("totalConnections", poolBean.getTotalConnections());
            stats.put("threadsAwaitingConnection", poolBean.getThreadsAwaitingConnection());
            stats.put("maximumPoolSize", configManager.getPool().getMaximumPoolSize());
            stats.put("minimumIdle", configManager.getPool().getMinimumIdle());
            stats.put("connectionTimeout", configManager.getPool().getConnectionTimeout());
            stats.put("idleTimeout", configManager.getPool().getIdleTimeout());
            stats.put("maxLifetime", configManager.getPool().getMaxLifetime());
            stats.put("initialized", initialized.get());
            stats.put("closed", dataSource.isClosed());

            return stats;
        } catch (Exception e) {
            log.error("获取TiDB连接池统计信息失败", e);
            return Collections.emptyMap();
        }
    }

    /**
     * 检查连接池健康状态
     */
    public boolean isHealthy() {
        if (!initialized.get() || dataSource == null || dataSource.isClosed()) {
            return false;
        }

        try (Connection conn = getConnection()) {
            return conn.prepareStatement(configManager.getSecurity().getConnectionTestQuery())
                      .executeQuery().next();
        } catch (Exception e) {
            log.error("TiDB连接池健康检查失败", e);
            return false;
        }
    }

    /**
     * 重新初始化连接池
     */
    public synchronized void reinitialize() {
        log.info("开始重新初始化TiDB连接池");
        initialized.set(false);
        if (dataSource != null && !dataSource.isClosed()) {
            dataSource.close();
        }
        initializeDataSource();
        log.info("TiDB连接池重新初始化完成");
    }

    /**
     * 获取配置管理器（用于其他组件获取配置信息）
     */
    public TiDBConfigManager getConfigManager() {
        return configManager;
    }

    /**
     * 关闭连接池
     */
    @PreDestroy
    public void shutdown() {
        if (dataSource != null && !dataSource.isClosed()) {
            log.info("正在关闭TiDB连接池...");
            dataSource.close();
            initialized.set(false);
            log.info("TiDB连接池已关闭");
        }
    }

    /**
     * 获取初始化状态
     */
    public boolean isInitialized() {
        return initialized.get();
    }
}
